# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin

from legged_gym import LEGGED_GYM_ROOT_DIR
import os
from os.path import join, dirname, abspath

import isaacgym
from legged_gym.envs import *
from legged_gym.utils import get_args, export_policy_as_jit, task_registry, Logger
from legged_gym.utils.diff_quat import flip_quat_by_w, quat_to_vec6d, quat_multiply, quat_inv, broadcast_quat_apply, broadcast_quat_multiply, quat_to_matrix, vec6d_to_quat

import numpy as np
import pickle
from transforms3d.axangles import axangle2mat
from transforms3d.quaternions import mat2quat
import torch
import time
from tqdm import tqdm


def get_root_states(env):
    root_states = torch.cat([
        (env.root_pos - env.env_origins) * env.obs_scales.root_pos,
        env.root_ori * env.obs_scales.root_ori,
        env.root_vel * env.obs_scales.root_vel,
        env.root_ang_vel * env.obs_scales.root_ang_vel,
    ], dim=1)
    assert root_states.shape[1] == 13
    return root_states


def get_698D_obs_np(env):
    obs_list = []

    obs_list.append(env.dof_pos * env.obs_scales.dof_pos)
    obs_list.append(env.dof_vel * env.obs_scales.dof_vel)

    obs_list.append(env.body_pos * env.obs_scales.body_pos)
    obs_list.append(env.body_ori * env.obs_scales.body_ori)
    obs_list.append(env.body_vel * env.obs_scales.body_vel)
    obs_list.append(env.body_ang_vel * env.obs_scales.body_ang_vel)

    obs_list.append(env.actions * env.obs_scales.last_action)

    obs_list.append(env.delayed_obs_target_jt_pos * env.obs_scales.target_dof_pos)
    obs_list.append(env.delayed_obs_target_jt_vel * env.obs_scales.target_dof_vel)

    pelvis_pos = env.body_pos_raw[:, :1, :] - env.env_origins.reshape(env.num_envs, 1, 3)  # shape = (num_envs, 1, 3)
    pelvis_ori = env.body_ori_raw[:, :1, :]  # shape = (num_envs, 1, 4)
    pelvis_ori_inv = quat_inv(pelvis_ori)
    target_body_pos = broadcast_quat_apply(pelvis_ori_inv, env.delayed_obs_target_body_pos.reshape(env.num_envs, (env.num_dofs + 1), 3) - pelvis_pos).reshape(-1, 3 * (env.num_dofs + 1))
    target_body_ori = flip_quat_by_w(broadcast_quat_multiply(pelvis_ori_inv, vec6d_to_quat(env.delayed_obs_target_body_ori.reshape(env.num_envs, (env.num_dofs + 1), 3, 2))))
    target_body_ori = quat_to_vec6d(target_body_ori).reshape(-1, 6 * (env.num_dofs + 1))
    target_body_vel = broadcast_quat_apply(pelvis_ori_inv, env.delayed_obs_target_body_vel.reshape(env.num_envs, (env.num_dofs + 1), 3)).reshape(-1, 3 * (env.num_dofs + 1))
    target_body_ang_vel = broadcast_quat_apply(pelvis_ori_inv, env.delayed_obs_target_body_ang_vel.reshape(env.num_envs, (env.num_dofs + 1), 3)).reshape(-1, 3 * (env.num_dofs + 1))
    obs_list.append(target_body_pos * env.obs_scales.target_body_pos)
    obs_list.append(target_body_ori * env.obs_scales.target_body_ori)
    obs_list.append(target_body_vel * env.obs_scales.target_body_vel)
    obs_list.append(target_body_ang_vel * env.obs_scales.target_body_ang_vel)
    
    obs_list.append(env.projected_gravity * env.obs_scales.base_ang_vel)
        
    if (hasattr(env.cfg.env.obs, "target_hand_positions")) and (env.cfg.env.obs.target_hand_positions):  # 6
        pelvis_pos = env.body_pos_raw[:, :1, :] - env.env_origins.reshape(env.num_envs, 1, 3)  # shape = (num_envs, 1, 3)
        pelvis_ori = env.body_ori_raw[:, :1, :]  # shape = (num_envs, 1, 4)
        pelvis_ori_inv = quat_inv(pelvis_ori)
        target_hand_positions = broadcast_quat_apply(pelvis_ori_inv, env.target_hand_positions.reshape(env.num_envs, 2, 3) - pelvis_pos).reshape(-1, 6)
        obs_list.append(target_hand_positions)
 
    obs_buf_np = torch.cat(obs_list, dim=-1).detach().cpu().numpy()
    return obs_buf_np


def inference_and_save(args):
    # env_cfg, train_cfg = task_registry.get_cfgs(name=args.task, load_run=args.load_run)
    env_cfg, train_cfg = task_registry.get_cfgs(name=args.task, load_run=args.run_name)
    # override some parameters for testing
    env_cfg.env.num_envs = min(env_cfg.env.num_envs, 400)
    env_cfg.terrain.num_rows = 5
    env_cfg.terrain.num_cols = 5
    env_cfg.terrain.curriculum = False
    env_cfg.noise.add_noise = False
    env_cfg.domain_rand.randomize_friction = False
    env_cfg.domain_rand.push_robots = False
    env_cfg.human.need_testset = True  # training set + test set

    # track tracked motion (check)
    if args.track_tracked_motion:
        env_cfg.human.need_testset = False
        env_cfg.human.file_dir = None
        env_cfg.human.track_tracked_motion_path = join(dirname(abspath(__file__)), "../../tracking_results_roam_10joint.pkl")

    # prepare environment
    env, _ = task_registry.make_env(name=args.task, args=args, env_cfg=env_cfg)

    # load policy
    train_cfg.runner.resume = True
    train_cfg.runner.run_name = 'play'

    ppo_runner, train_cfg = task_registry.make_alg_runner(env=env, name=args.task, args=args, train_cfg=train_cfg)
    policy = ppo_runner.get_inference_policy(device=env.device)

    obs = env.get_observations()
    root_states = get_root_states(env)
    
    # export policy as a jit module (used to run it from C++)
    if EXPORT_POLICY:
        path = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name, 'exported', 'policies')
        export_policy_as_jit(ppo_runner.alg.actor_critic, path)
        print('Exported policy as jit script to: ', path)

    logger = Logger(env.dt)
    robot_index = 0 # which robot is used for logging
    joint_index = 1 # which joint is used for logging
    stop_state_log = 10000 # number of steps before plotting states
    stop_rew_log = env.max_episode_length + 1 # number of steps before print average episode rewards

    # init inference results
    num_envs = env_cfg.env.num_envs  # 400
    num_obs = (698 + (6 if (hasattr(env.cfg.env.obs, "target_hand_positions")) and (env.cfg.env.obs.target_hand_positions) else 0)) if env.num_obs < 200 else env.num_obs  # 701
    num_actions = env.num_actions  # 19
    inference_results = [{
        "global_information_tracking": env.cfg.env.obs.root_pos,  # global information tracking or proprioception tracking
        "state_dimension": num_obs,  # 701
        "action_dimension": num_actions,  # 19
        "scene_mesh": {"vertices": None, "faces": None},  # vertices: (N_point, 3), faces: (N_face, 3)
        "scene_mesh_path": "",  # str
        "h1_start_root_global_translation": None,   # shape = (3,)
        "h1_start_root_global_rotation": None,  # shape = (4,), quaternion: (x, y, z, w)
        "root_states": [],  # shape = (N_frame, 13)
        "tracker_inputs": [],  # shape = (N_frame, num_obs)
        "tracked_actions": [],  # shape = (N_frame, num_actions)
        "tracked_torques": [],  # shape = (N_frame, num_actions)
        "tracked_dof_vels": [],  # shape = (N_frame, num_actions)
        "tracked_next_states": [],  # shape = (N_frame, num_obs)
        "tracked_next_root_states": [],  # shape = (N_frame, 13)
        "tracked_joint_angles": [],  # shape = (N_frame, num_actions)
    } for _ in range(num_envs)]

    # save scene mesh, h1 start root pose
    for i in range(num_envs):
        inference_results[i]["scene_mesh_path"] = env.scene_mesh_paths[env.env_scene_ids[i]]
        scene_mesh = env.scene_list[env.env_scene_ids[i]]
        inference_results[i]["scene_mesh"]["vertices"] = scene_mesh["vertices"].copy()
        inference_results[i]["scene_mesh"]["faces"] = scene_mesh["faces"].copy()
        start_root_global_translation = env.base_init_state[:3].clone()
        start_root_global_translation[:2] += env.target_all_info_seq["target_root_translation"][env.target_jt_i[i], 0, :2]
        inference_results[i]["h1_start_root_global_translation"] = start_root_global_translation.detach().cpu().numpy()
        
    termination = np.zeros(num_envs).astype(np.bool_)  # 每个环境是否已经结束rollout
    for frame_idx in tqdm(range(int(env.max_episode_length))):

        tracker_input_np = obs.clone().detach().cpu().numpy()  # shape = (num_envs, 701)
        root_states_np = root_states.clone().detach().cpu().numpy()  # shape = (num_envs, 13)
        # assert tracker_input_np.shape == (num_envs, num_obs)
        assert root_states_np.shape == (num_envs, 13)

        if env.cfg.env.policy_name == "ActorCriticPointNet":
            pointnet_feat, _, _ = ppo_runner.alg.pointnet(env.scene_pcd.permute(0, 2, 1).contiguous())
            obs = torch.cat([obs, pointnet_feat], dim=-1)

        actions = policy(obs.detach())
        action_bound = env.cfg.normalization.clip_actions
        clipped_actions = torch.clip(actions, -action_bound, action_bound).to(actions.device)
        actions = clipped_actions.clone()
        print("dof pos:", obs[:, 0:19].shape, obs[0, 0:19])
        print("actions:", actions.shape, actions[0])
        obs, _, rews, dones, infos = env.step(actions.detach())
        root_states = get_root_states(env)
        torques_np = env.torques.detach().cpu().numpy()
        dof_vel_np = env.dof_vel.detach().cpu().numpy()

        # time.sleep(0.2)

        # (for touch point) save task target hand positions
        if hasattr(env, "target_hand_positions"):
            task_target_hand_positions = env.target_hand_positions.detach().cpu().numpy()
            for i in range(num_envs):
                if not termination[i]:
                    inference_results[i]["task_target_hand_positions"] = task_target_hand_positions[i].copy()
        if hasattr(env, "kinematic_target_hand_positions"):
            kinematic_task_target_hand_positions = env.kinematic_target_hand_positions.detach().cpu().numpy()
            for i in range(num_envs):
                if not termination[i]:
                    inference_results[i]["kinematic_task_target_hand_positions"] = kinematic_task_target_hand_positions[i].copy()
        
        # check termination
        termination |= (dones > 0).detach().cpu().numpy()

        # save tracker input
        for i in range(num_envs):
            if not termination[i]:
                inference_results[i]["tracker_inputs"].append(tracker_input_np[i].copy())
                inference_results[i]["root_states"].append(root_states_np[i].copy())

        # save tracked_action
        actions_np = actions.detach().cpu().numpy()  # shape = (num_envs, 19)
        assert actions_np.shape == (num_envs, num_actions)
        for i in range(num_envs):
            if not termination[i]:
                inference_results[i]["tracked_actions"].append(actions_np[i].copy())
                inference_results[i]["tracked_torques"].append(torques_np[i].copy())
                inference_results[i]["tracked_dof_vels"].append(dof_vel_np[i].copy())

        # save tracked next state, joint global translations, joint global rotations
        next_obs_np = obs.detach().cpu().numpy()  # shape = (num_envs, 701)
        if next_obs_np.shape[1] < 200:  # for original HST
            next_obs_np = get_698D_obs_np(env)
        next_root_states_np = root_states.detach().cpu().numpy()  # shape = (num_envs, 13)
        assert next_obs_np.shape == (num_envs, num_obs)
        for i in range(num_envs):
            if not termination[i]:
                inference_results[i]["tracked_next_states"].append(next_obs_np[i].copy())
                inference_results[i]["tracked_next_root_states"].append(next_root_states_np[i].copy())
                inference_results[i]["tracked_joint_angles"].append(next_obs_np[i, 0:num_actions].copy())

        # print("step: {} / {}".format(i, int(env.max_episode_length)))

        if i < stop_state_log:
            logger.log_states(
                {
                    'dof_pos_target': actions[robot_index, joint_index].item() * env.cfg.control.action_scale,
                    'dof_pos': env.dof_pos[robot_index, joint_index].item(),
                    'dof_vel': env.dof_vel[robot_index, joint_index].item(),
                    'dof_torque': env.torques[robot_index, joint_index].item(),
                    'command_x': env.commands[robot_index, 0].item(),
                    'command_y': env.commands[robot_index, 1].item(),
                    'command_yaw': env.commands[robot_index, 2].item(),
                    'base_vel_x': env.base_lin_vel[robot_index, 0].item(),
                    'base_vel_y': env.base_lin_vel[robot_index, 1].item(),
                    'base_vel_z': env.base_lin_vel[robot_index, 2].item(),
                    'base_vel_yaw': env.base_ang_vel[robot_index, 2].item(),
                    'contact_forces_z': env.contact_forces[robot_index, env.feet_indices, 2].cpu().numpy()
                }
            )
        elif i==stop_state_log:
            logger.plot_states()
        if  0 < i < stop_rew_log:
            if infos["episode"]:
                num_episodes = torch.sum(env.reset_buf).item()
                if num_episodes>0:
                    logger.log_rewards(infos["episode"], num_episodes)
        elif i==stop_rew_log:
            logger.print_rewards()
    
    # statistics and change to np
    for i in range(num_envs):
        for key in inference_results[i]:
            if isinstance(inference_results[i][key], list):
                inference_results[i][key] = np.float32(inference_results[i][key])
                print("[sequence {}, key {}] shape =".format(i, key), inference_results[i][key].shape)

    # save to file
    save_dir = args.inference_save_dir
    save_filename = args.save_filename
    os.makedirs(save_dir, exist_ok=True)
    pickle.dump(inference_results, open(join(save_dir, save_filename + ".pkl"), "wb"))

if __name__ == '__main__':
    EXPORT_POLICY = True
    args = get_args(test=True)
    inference_and_save(args)
